package com.xxx.flink.transform;

import com.xxx.flink.customsource.CustomSource;
import com.xxx.flink.pojo.Event;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;

/**
 * 转换算子分区
 */
public class TransformPartitionTest {

    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        // 读取数据
        DataStreamSource<Event> sourceStream = env.addSource(new CustomSource());

        // 1.轮询分区（默认）
//        sourceStream.rebalance().print().setParallelism(4);

        // 2.随机分区shuffle（全部洗牌）
//        sourceStream.shuffle().print().setParallelism(4);

        // 3.重缩放分区recale（先分组，然后在组内轮询）
        DataStreamSource<Integer> streamSource = env.addSource(new RichParallelSourceFunction<Integer>() {
            @Override
            public void run(SourceContext<Integer> sourceContext) throws Exception {
                for (int i = 1; i <= 8; i++) {
                    // 将奇偶数分别发送到0号和1号分区
                    if (i % 2 == getRuntimeContext().getIndexOfThisSubtask()) {
                        sourceContext.collect(i);
                    }
                }
            }

            @Override
            public void cancel() {

            }
        }).setParallelism(2);
//        streamSource.rescale().print().setParallelism(4);
        // 测试结果：奇数1、3、5、7跑到3和4分区了，偶数2、4、6、8跑到1和2分区了

        // 4.广播分区
//        streamSource.broadcast().print().setParallelism(4);

        // 5.全局分区
        streamSource.global().print().setParallelism(4);

        env.execute();
    }
}
